package com.lili.other;

import java.util.*;

/**
 * @Auther: 李 力
 * @Date: 2024/8/6
 * @Description: com.lili.greedy
 * @version: 1.0
 */
public class Huffman {

    private static class Node {
        //字符
        Character ch;
        //频次
        int freq;
        Node left;

        @Override
        public String toString() {
            return "Node{" +
                    "ch=" + ch +
                    ", freq=" + freq +
                    '}';
        }

        Node right;
        //编码
        String code;

        public Node(int freq, Node left, Node right) {
            this.freq = freq;
            this.left = left;
            this.right = right;
        }

        public Node(String code) {
            this.code = code;
        }

        public Node(Character ch, int freq, Node left, Node right, String code) {
            this.ch = ch;
            this.freq = freq;
            this.left = left;
            this.right = right;
            this.code = code;
        }

        public Character getCh() {
            return ch;
        }

        public void setCh(Character ch) {
            this.ch = ch;
        }

        public int getFreq() {
            return freq;
        }

        public void setFreq(int freq) {
            this.freq = freq;
        }

        public Node getLeft() {
            return left;
        }

        public void setLeft(Node left) {
            this.left = left;
        }

        public Node getRight() {
            return right;
        }

        public void setRight(Node right) {
            this.right = right;
        }

        public String getCode() {
            return code;
        }

        public void setCode(String code) {
            this.code = code;
        }

        public Node(Character ch) {
            this.ch = ch;
        }

        public boolean isLeaf() {
            return left == null;
        }
    }

    private String str;

    Map<Character, Node> map = new HashMap<>();

    private Node root;

    public Huffman(String str) {
        this.str = str;
        //统计频次
        char[] charArray = str.toCharArray();
        for (char c : charArray) {
            if (!map.containsKey(c)) {
                Node node = new Node(c);
                map.put(c, node);
            }
            Node node = map.get(c);
            node.freq++;
        }
        //构造树
        PriorityQueue<Node> queue = new PriorityQueue<>(Comparator.comparingInt(Node::getFreq));
        queue.addAll(map.values());


        while (queue.size() >= 2) {
            Node x = queue.poll();
            Node y = queue.poll();
            int freq = x.freq + y.freq;
            Node node = new Node(freq, x, y);
            queue.offer(node);
        }

        Node root = queue.poll();
        this.root = root;
        //计算每个字符的编码
        dfs(root, new StringBuilder());
    }

    private int dfs(Node node, StringBuilder code) {
        int num = 0;//统计比特位
        if (node.isLeaf()) {
            //找到编码
            System.out.println(node + ":" + code);
            node.code = code.toString();
            num = node.freq * code.length();
        } else {
            num = num + dfs(node.left, code.append("0"));
            num = num + dfs(node.right, code.append("1"));
        }
        if (code.length() > 0) {
            code.deleteCharAt(code.length() - 1);
        }
        return num;
    }


    //编码
    public String encode() {
        char[] charArray = str.toCharArray();
        StringBuilder s = new StringBuilder();
        for (char c : charArray) {
            Node node = map.get(c);
            String code = node.code;
            s.append(code);
        }
        return s.toString();
    }

    //解码
    public String decode(String str) {
        char[] charArray = str.toCharArray();
        int i = 0;
        StringBuilder s = new StringBuilder();
        Node node = root;
        while (i < charArray.length) {
            if (!node.isLeaf()) {//非叶子节点
                if (charArray[i] == '0') {
                    //向左走
                    node = node.left;
                } else if (charArray[i] == '1') {
                    //向右走
                    node = node.right;
                }
                i++;
            }
            if (node.isLeaf()) {
                //叶子节点
                s.append(node.ch);
                node = root;
            }
        }
        return s.toString();
    }

}
